package com.qiuzhao.config;

import ch.qos.logback.core.util.StringUtil;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.qiuzhao.entity.ChatMemory;
import com.qiuzhao.mapper.ChatMemoryMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Repository;

import java.util.Date;
import java.util.List;
import java.util.stream.Collectors;

@Repository
@Slf4j
public class MySqlChatMemoryRepository implements ChatMemoryRepository {
    @Autowired
    private ChatMemoryMapper chatMemoryMapper;
    @Autowired
    private ChatModel chatModel;

    @Override
    public List<String> findConversationIds() {
        return chatMemoryMapper.selectList(null).stream().map(ChatMemory::getConversationId).collect(Collectors.toList());
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        if (StringUtil.isNullOrEmpty(conversationId))
            return null;
        LambdaQueryWrapper<ChatMemory> lambdaQueryWrapper = new LambdaQueryWrapper<>();
        lambdaQueryWrapper.eq(ChatMemory::getConversationId, conversationId);
        List<ChatMemory> chatMemories = chatMemoryMapper.selectList(lambdaQueryWrapper);
        return chatMemories.stream().map(chatMemory -> {
            String role = chatMemory.getRole();
            String content = chatMemory.getContent();
            switch (role.toLowerCase()) {
                case "user":
                    return new UserMessage(content);
                case "assistant":
                    return new AssistantMessage(content);
                case "system":
                    return new SystemMessage(content);
                default:
                    throw new IllegalArgumentException("未知角色: " + role);
            }
        }).collect(Collectors.toList());
    }

    @Override
    public void saveAll(String conversationId, List<Message> messages) {
        if (messages.isEmpty()){
            return;
        }
        Message lastMessage = messages.get(messages.size() - 1);
        String[] split = conversationId.split("-");
        conversationId = split[0];
        String roleId = split[1];
        List<ChatMemory> memory = chatMemoryMapper.selectList(new LambdaQueryWrapper<ChatMemory>().eq(ChatMemory::getConversationId, conversationId));
        String title;
        //如果是第一次生成，那么需要生成一个标题
        if (memory == null || memory.isEmpty()){
            log.info(lastMessage.getText());
            title = chatModel.call("请根据以下内容，生成一个会话标题，20字以内。" + lastMessage.getText());
        }else {
            title = memory.get(0).getTitle();
        }
        ChatMemory chatMemory = new ChatMemory();
        chatMemory.setConversationId(conversationId);
        chatMemory.setSystemRole(roleId);
        chatMemory.setRole(lastMessage.getMessageType().getValue());
        chatMemory.setContent(lastMessage.getText());
        chatMemory.setTitle(title);
        chatMemory.setCreateTime(new Date());
        chatMemoryMapper.insert(chatMemory);
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        LambdaQueryWrapper<ChatMemory> lambdaQueryWrapper = new LambdaQueryWrapper<ChatMemory>()
                .eq(ChatMemory::getConversationId, conversationId);
        chatMemoryMapper.delete(lambdaQueryWrapper);
    }
}
